import tensorflow as tf

from .registry import register
from .defaults import default, default_cifar10


# from https://github.com/tensorflow/models/blob/master/resnet/resnet_main.py
@register
def resnet_default():
  resnet_default = default_cifar10()
  resnet_default.model = "resnet"
  resnet_default.batch_size = 128
  resnet_default.learning_rate = 0.1
  resnet_default.lr_scheme = "resnet"
  resnet_default.relu_leakiness = 0.1
  resnet_default.use_bottleneck = False
  resnet_default.num_residual_units = 5
  resnet_default.weight_decay_rate = 2e-4
  resnet_default.num_classes = 10
  resnet_default.optimizer = "momentum"
  return resnet_default


@register
def resnet_cifar100():
  resnet_cifar100 = resnet_default()
  resnet_cifar100.num_classes = 100
  resnet_cifar100.output_shape = [100]
  return resnet_cifar100


@register
def cifar10_resnet32():
  hps = resnet_default()
  hps.initializer = "glorot_normal_initializer"

  return hps


@register
def cifar10_resnet32_no_dropout():
  hps = cifar10_resnet32()
  hps.drop_rate = 0.0

  return hps


@register
def cifar10_resnet32_trgtd_weight():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_weight"
  hps.targ_rate = 0.5

  return hps


@register
def cifar10_resnet32_weight():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.25
  hps.dropout_type = "untargeted_weight"

  return hps


@register
def cifar10_resnet32_weight_50():
  hps = cifar10_resnet32_weight()
  hps.drop_rate = 0.50

  return hps


@register
def cifar10_resnet32_trgtd_unit():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_unit"
  hps.targ_rate = 0.5

  return hps


@register
def cifar10_resnet32_trgtd_ard():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.25
  hps.dropout_type = "targeted_ard"
  hps.targ_rate = 0.5

  return hps


@register
def cifar10_resnet32_unit():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.25
  hps.dropout_type = "untargeted_unit"

  return hps


@register
def cifar10_resnet32_unit_50():
  hps = cifar10_resnet32_unit()
  hps.drop_rate = 0.50

  return hps


@register
def cifar10_resnet32_l1():
  hps = cifar10_resnet32_no_dropout()
  hps.l1_norm = 0.001

  return hps


@register
def cifar10_resnet32_l1_01():
  hps = cifar10_resnet32_no_dropout()
  hps.l1_norm = 0.01

  return hps


@register
def cifar10_resnet32_l1_1():
  hps = cifar10_resnet32_no_dropout()
  hps.l1_norm = 0.1

  return hps


@register
def cifar10_resnet32_trgted_weight_l1():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_weight"
  hps.targ_rate = 0.5
  hps.l1_norm = 0.1

  return hps


@register
def cifar10_resnet32_targeted_unit_l1():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_unit"
  hps.targ_rate = 0.5
  hps.l1_norm = 0.1

  return hps


@register
def cifar10_resnet32_trgtd_unit_botk75_33():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.33
  hps.dropout_type = "targeted_unit"
  hps.targ_rate = 0.75

  return hps

@register
def cifar10_resnet32_trgtd_unit_botk75_50():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_unit"
  hps.targ_rate = 0.75

  return hps

@register
def cifar10_resnet32_trgtd_unit_botk75_66():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.66
  hps.dropout_type = "targeted_unit"
  hps.targ_rate = 0.75

  return hps


@register
def cifar10_resnet32_trgtd_weight_botk75_33():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.33
  hps.dropout_type = "targeted_weight"
  hps.targ_rate = 0.75

  return hps


@register
def cifar10_resnet32_trgtd_weight_botk75_66():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.66
  hps.dropout_type = "targeted_weight"
  hps.targ_rate = 0.75

  return hps


@register
def cifar10_resnet32_louizos():
  hps = cifar10_resnet32_no_dropout()
  hps.louizos_beta = 2. / 3.
  hps.louizos_zeta = 1.1
  hps.louizos_gamma = -0.1
  hps.louizos_cost = 0.001
  hps.dropout_type = "louizos_weight"
  hps.drop_rate = 0.25

  return hps


@register
def cifar10_resnet32_louizos_weight_1():
  hps = cifar10_resnet32_louizos()
  hps.louizos_cost = 0.1
  hps.dropout_type = "louizos_weight"
  hps.drop_rate = 0.25

  return hps


@register
def cifar10_resnet32_louizos_weight_2():
  hps = cifar10_resnet32_louizos()
  hps.louizos_cost = 0.01

  return hps


@register
def cifar10_resnet32_louizos_weight_3():
  hps = cifar10_resnet32_louizos()
  hps.louizos_cost = 0.005

  return hps


@register
def cifar10_resnet32_louizos_weight_4():
  hps = cifar10_resnet32_louizos()
  hps.louizos_cost = 0.0001

  return hps


@register
def cifar10_resnet32_louizos_unit():
  hps = cifar10_resnet32_no_dropout()
  hps.louizos_beta = 2. / 3.
  hps.louizos_zeta = 1.1
  hps.louizos_gamma = -0.1
  hps.louizos_cost = 0.001
  hps.dropout_type = "louizos_unit"
  hps.drop_rate = 0.25

  return hps


@register
def cifar10_resnet32_louizos_unit_1():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.1

  return hps


@register
def cifar10_resnet32_louizos_unit_2():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.01

  return hps


@register
def cifar10_resnet32_louizos_unit_3():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.005

  return hps


@register
def cifar10_resnet32_louizos_unit_4():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.0001

  return hps


@register
def cifar10_resnet32_louizos_unit_5():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.00001

  return hps


@register
def cifar10_resnet32_louizos_unit_6():
  hps = cifar10_resnet32_louizos_unit()
  hps.louizos_cost = 0.000001

  return hps


@register
def cifar10_resnet32_variational():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "variational"
  hps.drop_rate = 0.75
  hps.thresh = 3
  hps.var_scale = 1. / 100

  return hps


@register
def cifar10_resnet32_variational_unscaled():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "variational"
  hps.drop_rate = 0.75
  hps.thresh = 3
  hps.var_scale = 1

  return hps


@register
def cifar10_resnet32_variational_unit():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "variational_unit"
  hps.drop_rate = 0.75
  hps.thresh = 3
  hps.var_scale = 1. / 100

  return hps


@register
def cifar10_resnet32_variational_unit_unscaled():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "variational_unit"
  hps.drop_rate = 0.75
  hps.thresh = 3
  hps.var_scale = 1

  return hps


@register
def cifar10_resnet32_smallify_1eneg4():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "smallify_dropout"
  hps.smallify = 1e-4
  hps.smallify_mv = 0.9
  hps.smallify_thresh = 0.5

  return hps


@register
def cifar10_resnet32_smallify_1eneg3():
  hps = cifar10_resnet32_smallify_1eneg4()
  hps.smallify = 1e-3

  return hps


@register
def cifar10_resnet32_smallify_1eneg5():
  hps = cifar10_resnet32_smallify_1eneg4()
  hps.smallify = 1e-5

  return hps


@register
def cifar10_resnet32_smallify_1eneg6():
  hps = cifar10_resnet32_smallify_1eneg4()
  hps.smallify = 1e-6

  return hps


@register
def cifar10_resnet32_smallify_weight_1eneg4():
  hps = cifar10_resnet32_no_dropout()
  hps.dropout_type = "smallify_weight_dropout"
  hps.smallify = 1e-4
  hps.smallify_mv = 0.9
  hps.smallify_thresh = 0.5

  return hps


@register
def cifar10_resnet32_smallify_weight_1eneg3():
  hps = cifar10_resnet32_smallify_weight_1eneg4()
  hps.smallify = 1e-3

  return hps


@register
def cifar10_resnet32_smallify_weight_1eneg5():
  hps = cifar10_resnet32_smallify_weight_1eneg3()
  hps.smallify = 1e-5

  return hps


@register
def cifar10_resnet32_smallify_weight_1eneg6():
  hps = cifar10_resnet32_smallify_weight_1eneg3()
  hps.smallify = 1e-6

  return hps


@register
def cifar10_resnet32_trgtd_weight_random():
  hps = cifar10_resnet32_no_dropout()
  hps.drop_rate = 0.5
  hps.dropout_type = "targeted_weight_random"
  hps.targ_rate = 0.5

  return hps


@register
def cifar10_resnet32_trgtd_unit_random():
  hps = cifar10_resnet32_trgtd_weight_random()
  hps.dropout_type = "targeted_unit_random"

  return hps


@register
def cifar10_resnet32_trgtd_weight_random_botk75_66():
  hps = cifar10_resnet32_trgtd_weight_random()
  hps.drop_rate = 0.66
  hps.dropout_botk = 0.75

  return hps


@register
def cifar10_resnet32_trgtd_unit_random_botk75_66():
  hps = cifar10_resnet32_trgtd_weight_random_botk75_66()
  hps.dropout_type = "targeted_unit_random"

  return hps


# ================================
